import subprocess
from pprint import pprint
from util import dig_dns
from multiprocessing.dummy import Pool
import re
import json
import sys

def my_dig_dns(ip):
    return dig_dns('@' + ip, 'www.baidu.com'):

def dig(domain):
    with open('dns.json','rt')as f:
        dns_list = json.load(f)
    ret_ips = []

    pool = Pool(100)
    ret_list = pool.map(my_dig_dns,dns_list)
    pool.close()
    pool.join()
    for data in ret_list:
        if data:
            for ip in data:
                if re.match('[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}',ip):
                    ret_ips.append(ip)
    ret_ips = set(ret_ips)
    return ret_ips

assert(len(sys.argv)<3)
if len(sys.argv) == 2:
    domain = sys.argv[1]
else:
    domain = 'www.baidu.com'

print('dig the domain: ' + domain)
ret = dig(domain)
print(ret)
